# -*- coding: utf-8 -*-
# May Saint Diana bless you!
# Created Time : 2021/11/1 15:45
# Author:Zhou
import torch
import torch.nn.functional as F  # 激励函数都在这
from torch.autograd import Variable
import matplotlib.pyplot as plt
import numpy as np
import random
import os
from ANC_tools import insertToTail, ANR, FIR_Identify

# 读取实验数据
# path = "次级通道建模数据.csv"  # 输入是传递给次级扬声器播放的信号（高斯白噪声，能激发出所有频段），输出是在误差传感器处采集到的误差信号
path = "非线性5.csv"  # 一个参考信号和该信号经过非线性过程的对应输出
data = np.loadtxt(path, delimiter=",")

# 构造简单训练数据 初始数据规模为
x = []
y = []

for i in range(1000):
    time_point = random.randint(31, 45000)  # 找一个随机时刻
    x.append(data[time_point - 31:time_point + 1, 0])  # 训练输入数据是当前及过去15个时刻的输出值
    y.append(data[time_point, 1])  # 输出是当前时刻的传感器数据  认为：当前的y和过去若干时刻的输入有关
# x y 放进numpy，再转成pytorch的tensor数据
x = np.array(x)
y = np.array(y)
x = Variable(torch.from_numpy(x)).to(torch.float)
y = Variable(torch.from_numpy(y[:, np.newaxis])).to(torch.float)
# 训练数据：1000组，每组输入是32*1，输出是1，即：构建当前时刻的误差信号和过去若干时刻（16）的次级扬声器输出信号之间的映射关系:
# y(n)=f[x(n),x(n-1),x(n-2),x(n-3),...,x(n-31)]


#  自定义一个全连接网络 32*64*16*1  当前的输出和之前32个时间点的输入有关
net = torch.nn.Sequential(
    torch.nn.Linear(32, 64),
    torch.nn.ReLU(),
    torch.nn.Linear(64, 16),
    torch.nn.ReLU(),
    torch.nn.Linear(16, 1))

optimizer = torch.optim.SGD(net.parameters(), lr=0.1)  # 用SGD优化器优化，lr为学习率，将net的参数传入
loss_func = torch.nn.MSELoss()

# 如果存在训练过的网络参数文件，直接读取
if os.path.exists('spm_net_params.pkl'):
    net.load_state_dict(torch.load('spm_net_params.pkl'))  # 提取net的参数

plt.ion()  # 画图
plt.show()
print("----------开始BP训练----------")
for t in range(1000):
    prediction = net(x)  # 通过网络计算得到的预测值  暂时没有用批训练，因此每次都是将1000个x先后输入网络，得到1000个输出y'
    loss = loss_func(prediction, y)  # 计算y'和y的误差，得到损失函数 0
    optimizer.zero_grad()  # 将net.parameters中上一批梯度清零
    loss.backward()  # 反向传递  计算节点梯度
    optimizer.step()  # 优化节点
    if t % 5 == 0:  # 动态显示
        # plot and show learning process
        plt.cla()
        plt.plot(y.data.numpy(), lw=5)
        plt.plot(prediction.data.numpy(), 'r-', lw=5)
        plt.text(25, 3, 'Loss=%.4f' % loss.data.numpy(), fontdict={'size': 20, 'color': 'green'})
        plt.pause(0.1)
print("----------训练完成！----------")

# 保存网络的参数
torch.save(net.state_dict(), 'spm_net_params.pkl')  # 只保存网络中的参数 (速度快, 占内存少)

plt.ioff()
plt.show()

# 验证拟合效果
signal_x = data[:, 0]
signal_y = data[:, 1]
cx = np.zeros(32)
BP_error = np.zeros(len(signal_x))
print("----------验证拟合效果----------")
for i in range(len(signal_x)):
    cx = insertToTail(cx, signal_x[i])  # 需要和网络的数据对应 从上至下应该是 x(n-15) 到 x(n)，所以要插入到末尾
    BP_x = Variable(torch.from_numpy(cx)).to(torch.float)
    BP_pred = net(BP_x)
    BP_error[i] = BP_pred.data.numpy() - signal_y[i]
print("平均误差：{0},最大误差：{1}".format(np.average(np.abs(BP_error)),
                                 np.max(np.abs(BP_error))))

# 用FIR滤波器对次级路径进行建模，得到滤波系数和拟合误差
(FIR_w, FIR_error) = FIR_Identify(signal_x, signal_y, 32, 0.0001)

# 比较两种方式的效果
plt.figure()
plt.subplot(221)
plt.plot(signal_y)
plt.plot(FIR_error)
plt.legend(['y', 'error'])
plt.title('FIR')
plt.subplot(222)
plt.plot(ANR(signal_y, FIR_error))
plt.title('FIR_ANR')

plt.subplot(223)
plt.plot(signal_y)
plt.plot(BP_error)
plt.legend(['y', 'error'])
plt.title('BP')
plt.subplot(224)
plt.plot(ANR(signal_y, BP_error))
plt.title('BP_ANR')
# 展示
plt.show()
